import numpy as np
import matplotlib.pyplot as plt
from skimage import io, filters, feature
from scipy.ndimage import uniform_filter
import joblib

# 读取保存的模型
clf = joblib.load('random_forest_model.joblib')

# 数据准备和预处理
def load_and_preprocess_images(image_path, mask_path):
    image = io.imread(image_path)
    mask = io.imread(mask_path)

    # 确保掩模是单通道
    if mask.ndim == 3:
        mask = mask[:, :, 0]

    # 检查掩模与图像大小一致
    if mask.shape != image.shape[:2]:
        raise ValueError("掩模形状 {} 与图像形状 {} 不匹配".format(mask.shape, image.shape))

    return image, mask

# 应用不同类型的滤波器
def extract_features(image):
    features = []

    # 均值滤波
    mean_filtered = uniform_filter(image, size=3)
    features.append(mean_filtered.flatten())

    # 高斯滤波
    gaussian_filtered = filters.gaussian(image, sigma=1)
    features.append(gaussian_filtered.flatten())

    # Sobel滤波
    sobel_filtered = filters.sobel(image)
    features.append(sobel_filtered.flatten())

    # Canny边缘检测
    canny_edges = feature.canny(image)
    features.append(canny_edges.flatten())

    return np.array(features).T

# 加载和预处理数据
image_2, mask_2 = load_and_preprocess_images("E:\\git_project\\qaz\\Sandstone_2.tif",
                                              "E:\\git_project\\qaz\\Sandstone_2_segment.tif")

# 特征提取
X_2 = extract_features(image_2)
y_2 = mask_2.flatten()

# 使用模型进行预测
y_pred = clf.predict(X_2)

# 将预测结果转换为与掩模相同的形状
predicted_mask = y_pred.reshape(mask_2.shape)

# 计算准确率
accuracy = np.mean(predicted_mask == mask_2)
print(f"准确率: {accuracy:.3f}")

# 显示原图
plt.figure(figsize=(15, 5))
plt.subplot(1, 3, 1)
plt.imshow(image_2, cmap='gray')
plt.title('Original Image')

# 显示掩模
plt.subplot(1, 3, 2)
plt.imshow(mask_2, cmap='gray')
plt.title('Segment')

# 显示预测结果
plt.subplot(1, 3, 3)
plt.imshow(predicted_mask, cmap='gray')
plt.title(f'Segmentation (acc: {accuracy:.3f})')
plt.show()